%% copy riak-0.14.2 source to learn dynamo implementation and erlang
%% author: lichuang
%% copy source from riak_kv_vnode.erl

%% edynamo_kv_vnode: VNode Implementation
-module(edynamo_kv_vnode).
-behaviour(edynamo_vnode).

-include_lib("riak_kv_vnode.hrl").
-include_lib("riak_core_pb.hrl").
%% API
-export([start_vnode/1,
         get/3,
         mget/3,
         del/3,
         put/6,
         readrepair/6,
         list_keys/4,
         fold/3,
         get_vclocks/2]).

%% edynamo_vnode API
-export([init/1,
         terminate/2,
         handle_command/3,
         is_empty/1,
         delete/1,
         handle_handoff_command/3,
         handoff_starting/2,
         handoff_cancelled/1,
         handoff_finished/2,
         handle_handoff_data/2,
         encode_handoff_item/2]).

-record(state, {idx :: partition(),
                mod :: module(),
                modstate :: term(),
                mrjobs :: term(),
                in_handoff = false :: boolean()}).

-record(putargs, {returnbody :: boolean(),
                  lww :: boolean(),
                  bkey :: {binary(), binary()},
                  robj :: term(),
                  reqid :: non_neg_integer(),
                  bprops :: maybe_improper_list(),
                  starttime :: non_neg_integer(),
                  prunetime :: undefined| non_neg_integer()}).

%% API
start_vnode(I) ->
  edynamo_vnode_master:get_vnode_pid(I, edynamo_kv_vnode).

get(Preflist, BKey, ReqId) ->
  Req = ?KV_GET_REQ{bkey = BKey,
                    req_id = ReqId},
  %% Assuming this function is called from a FSM process
  %% so self() == FSM pid
  edynamo_vnode_master:sync_spawn_command(Preflist,
                                          Req,
                                          {fms, ReqId, self()},
                                          edynamo_vnode_master).

mget(Preflist, BKeys, ReqId) ->
  Req = ?KV_MGET_REQ{bkeys = BKeys,
                     req_id = ReqId,
                     from = {fsm, self()}},
  edynamo_vnode_master:command(Preflist,
                               Req,
                               edynamo_vnode_master).

del(Preflist, BKey, ReqId) ->
  edynamo_vnode_master:sync_command(Preflist,
                                    ?KV_DELETE_REQ{bkey = BKey,
                                                   req_id = ReqId},
                                    edynamo_vnode_master).
                                
%% Issue a put for the object to the preflist, expecting a reply
%% to an FSM.
put(Preflist, BKey, Obj, ReqId, StartTime, Options) when is_integer(StartTime) ->
  put(Preflist, BKey, Obj, ReqId, StartTime, Options, {fsm, undefined, self()}).

put(Preflist, BKey, Obj, ReqId, StartTime, Options, Sender)
  when is_integer(StartTime) ->
    edynamo_vnode_master:command(Preflist,
                                 ?KV_PUT_REQ{
                                  bkey = BKey,
                                  object = Obj,
                                  req_id = ReqId,
                                  start_time = StartTime,
                                  options = Options},
                                 Sender,
                                 edynamo_vnode_master).

%% Do a put without sending any replies
readrepair(Preflist, BKey, Obj, ReqId, StartTime, Options) ->
  put(Preflist, BKey, Obj, ReqId, StartTime, [rr | Options], ignore).

list_keys(Preflist, ReqId, Caller, Bucket) ->
  edynamo_vnode_master:command(Preflist,
                               ?KV_LISTKEYS_REQ{
                                bucket = Bucket,
                                req_id = ReqId,
                                caller = Caller},
                              ignore,
                              edynamo_vnode_master).

fold(Preflist, Fun, Acc0) ->
  edynamo_vnode_master:sync_spawn_command(Preflist,
                                          ?FOLD_REQ{
                                            foldfun = Fun,
                                            acc0    = Acc0},
                                          edynamo_vnode_master).

get_vclocks(Preflist, BKeyList) ->
  edynamo_vnode_master:sync_spawn_command(Preflist,
                                          ?KV_VCLOCK_REQ{bkeys = BKeyList},
                                          edynamo_vnode_master).

%% VNode callbacks
init([Index]) ->
  Mod = app_helper:get_env(edynamo, storage_backend),
  Configuration = app_helper:get_env(edynamo),
  {ok, ModState} = Mod:start(Index, Configuration),

  {ok, #state{idx = Index, mod = Mod, modstate = ModState, mrjobs = dict:new()}}.

handle_command(?KV_PUT_REQ{bkey = BKey,
                           object = Object,
                           req_id = ReqId,
                           start_time = StartTime,
                           options = Options},
               Sender, State=#state{idx = Idx}) ->
  %%riak_kv_mapred_cache:eject(BKey),
  error_logger:error_msg("before put, key:~p~n", [BKey]),
  edynamo_vnode:reply(Sender, {w, Idx, ReqId}),
  do_put(Sender, BKey, Object, ReqId, StartTime, Options, State),
  {noreply, State};

handle_command(?KV_GET_REQ{bkey = BKey, req_id = ReqId}, Sender, State) ->
  do_get(Sender, BKey, ReqId, State);
handle_command(#riak_kv_listkeys_req_v1{bucket = Bucket, req_id = ReqId}, _Sender,
               State=#state{mod = Mod, modstate = ModState, idx = Idx}) ->
  do_list_bucket(ReqId, Bucket, Mod, ModState, Idx, State);
handle_command(?KV_LISTKEYS_REQ{bucket = Bucket, req_id = ReqId, caller = Caller}, _Sender,
               State = #state{mod = Mod, modstate = ModState, idx = Idx}) ->
  do_list_keys(Caller, ReqId, Bucket, Idx, Mod, ModState),
  {noreply, State};
handle_command(?KV_DELETE_REQ{bkey = BKey, req_id = ReqId}, _Sender,
               State = #state{mod = Mod, modstate = ModState,
                              idx = Idx}) ->
  %%riak_kv_mapred_cache:eject(BKey),
  case Mod:delete(ModState, BKey) of
    ok ->
      {reply, {del, Idx, ReqId}, State};
    {error, _Reason} ->
      {reply, {fail, Idx, ReqId}, State}
  end;
handle_command(?KV_VCLOCK_REQ{bkeys = BKeys}, _Sender, State) ->
  {reply, do_get_vclocks(BKeys, State), State};
handle_command(?FOLD_REQ{foldfun = Fun, acc0 = Acc}, _Sender, State) ->
  Reply = do_fold(Fun, Acc, State),
  {reply, Reply, State};

%% Commands originating from inside this vnode
handle_command({backend_callback, Ref, Msg}, _Sender,
               State = #state{mod = Mod, modstate = ModState}) ->
  Mod:callback(ModState, Ref, Msg),
  {noreply, State}.

handle_handoff_command(Req = ?FOLD_REQ{}, Sender, State) ->
  handle_command(Req, Sender, State);
handle_handoff_command(Req = {backend_callback, _Ref, _Msg}, Sender, State) ->
  handle_command(Req, Sender, State);
handle_handoff_command(_Req, _Sender, State) ->
  {forward, State}.

handoff_starting(_TargetNode, State) ->
  {true, State#state{in_handoff = true}}.

handoff_cancelled(State) ->
  {ok, State#state{in_handoff = false}}.

handoff_finished(_TargetNode, State) ->
  {ok, State}.

handle_handoff_data(BinObj, State) ->
  PBObj = edynamo_pb:decode_riakobject_pb(zlib:unzip(BinObj)),
  BKey = {PBObj#riakobject_pb.bucket, PBObj#riakobject_pb.key},
  case do_diffobj_put(BKey, binary_to_term(PBObj#riakobject_pb.val), State) of
    ok ->
      {reply, ok, State};
    Err ->
      {reply, {error, Err}, State}
  end.

encode_handoff_item({B, K}, V) ->
  zlib:zip(edynamo_pb:encode_riakobject_pb(
            #riakobject_pb{bucket = B, key = K, val = V})).

is_empty(State=#state{mod = Mod, modstate = ModState}) ->
  {Mod:is_empty(ModState), State}.

delete(State=#state{mod = Mod, modstate = ModState}) ->
  ok = Mod:drop(ModState),
  {ok, State}.

terminate(_Reason, #state{mod = Mod, modstate = ModState}) ->
  Mod:stop(ModState),
  ok.

%% @private
% upon receipt of a client-initiated put
do_put(Sender, {Bucket, _key} = BKey, RObj, ReqID, StartTime, Options, State) ->
  {ok, Ring} = edynamo_ring_manager:get_my_ring(),
  BProps = edynamo_bucket:get_bucket(Bucket, Ring),
  case proplists:get_value(rr, Options, false) of
    true ->
      PruneTime = undefined;
    false ->
      PruneTime = StartTime
  end,
  PutArgs = #putargs{returnbody = proplists:get_value(returnbody, Options, false),
                     lww = proplists:get_value(last_write_wins, BProps, false),
                     bkey = BKey,
                     robj = RObj,
                     reqid = ReqID,
                     bprops = BProps,
                     starttime = StartTime,
                     prunetime = PruneTime},
  Reply = perform_put(prepare_put(State, PutArgs), State, PutArgs),
  edynamo_vnode:reply(Sender, Reply).
  %%riak_kv_stat:update(vnode_put).
  
prepare_put(#state{}, #putargs{lww = true, robj = RObj}) ->
  {true, RObj};
prepare_put(#state{mod = Mod, modstate = ModState}, #putargs{bkey = BKey,
                                                             robj = RObj,
                                                             reqid = ReqID,
                                                             bprops = BProps,
                                                             starttime = StartTime,
                                                             prunetime = PruneTime}) ->
  case syntactic_put_merge(Mod, ModState, BKey, RObj, ReqID, StartTime) of
    {oldobj, OldObj} ->
      {false, OldObj};
    {newobj, NewObj} ->
      VC = edynamo_object:vclock(NewObj),
      AMObj = enforce_allow_mult(NewObj, BProps),
      case PruneTime of
        undefined ->
          ObjToStore = AMObj;
        _ ->
          ObjToStore = edynamo_object:set_vclock(AMObj,
                                                 vclock:prune(VC, PruneTime, BProps))
      end,
      {true, ObjToStore}
  end.

perform_put({false, Obj}, #state{idx = Idx}, #putargs{returnbody = true, reqid = ReqID}) ->
  {dw, Idx, Obj, ReqID};
perform_put({false, _Obj}, #state{idx = Idx}, #putargs{returnbody = false, reqid = ReqId}) ->
  {dw, Idx, ReqId};
perform_put({true, Obj}, #state{idx = Idx, mod = Mod, modstate = ModState},
            #putargs{returnbody = RB, bkey = BKey, reqid = ReqID}) ->
  Val = term_to_binary(Obj),
  case Mod:put(ModState, BKey, Val) of
    ok ->
      error_logger:error_msg("put ok, key:~p, val:~p~n", [BKey, binary_to_term(Val)]),
      case RB of
        true -> {dw, Idx, Obj, ReqID};
        false -> {dw, Idx, ReqID}
      end;
    {error, _Reason} ->
      {fail, Idx, ReqID}
  end.

%% @private
%% enforce allow_mult bucket property so that no backend ever stores
%% an object with multiple contents if allow_mult=false for that bucket
enforce_allow_mult(Obj, BProps) ->
  case proplists:get_value(allow_mult, BProps) of
    true -> Obj;
    _ ->
      case edynamo_object:get_contents(Obj) of
        [_] -> Obj;
        Mult ->
          {MD, V} = select_newest_content(Mult),
          edynamo_object:set_contents(Obj, [{MD, V}])
      end
  end.

%% @private
%% choose the latest content to store for the allow_mult=false case
select_newest_content(Mult) ->
  hd(lists:sort(
      fun({MD0, _}, {MD1, _}) ->
          edynamo_util:compare_dates(
            dict:fetch(<<"X-Riak-Last-Modified">>, MD0),
            dict:fetch(<<"X-Riak-Last-Modified">>, MD1))
      end,
      Mult)).

%% @private
syntactic_put_merge(Mod, ModState, BKey, Obj1, ReqId) ->
  syntactic_put_merge(Mod, ModState, BKey, Obj1, ReqId, vclock:timestamp()).

syntactic_put_merge(Mod, ModState, BKey, Obj1, ReqId, StartTime) ->
  case Mod:get(ModState, BKey) of
    {error, notfound} -> {newobj, Obj1};
    {ok, Val0} ->
      Obj0 = binary_to_term(Val0),
      ResObj = edynamo_object:syntactic_merge(
                 Obj0, Obj1, term_to_binary(ReqId), StartTime),
      case edynamo_object:vclock(ResObj) =:= edynamo_object:vclock(Obj0) of
        true -> {oldobj, ResObj};
        false -> {newobj, ResObj}
      end
  end.

%% @private
do_get(_Sender, BKey, ReqID,
       State=#state{idx = Idx, mod = Mod, modstate = ModState}) ->
  Retval = do_get_term(BKey, Mod, ModState),
  error_logger:error_msg("key:~p, ret:~p~n", [BKey, Retval]),
  %%riak_kv_stat:update(vnode_get),
  {reply, {r, Retval, Idx, ReqID}, State}.

%% @private
do_get_term(BKey, Mod, ModState) ->
  case do_get_binary(BKey, Mod, ModState) of
    {ok, Bin} ->
      {ok, binary_to_term(Bin)};
    Err ->
      Err
  end.

do_get_binary(BKey, Mod, ModState) ->
  Mod:get(ModState, BKey).

%% @private
do_list_bucket(ReqID, Bucket, Mod, ModState, Idx, State) ->
  Retval = Mod:list_bucket(ModState, Bucket),
  {reply, {kl, Retval, Idx, ReqID}, State}.

%% @private
do_list_keys(Caller, ReqId, Bucket, Idx, Mod, ModState) ->
  F = fun(BKey, _, Acc) ->
        process_keys(Caller, ReqId, Idx, Bucket, BKey, Acc) end,
  case Mod:fold(ModState, F, []) of
    [] ->
      ok;
    Remainder ->
      Caller ! {ReqId, {kl, Idx, Remainder}}
  end,
  Caller ! {ReqId, Idx, done}.

%% @private
process_keys(Caller, ReqId, Idx, '_', {Bucket, _K}, Acc) ->
  %% Bucket='_' means "list buckets" instead of "list keys"
  buffer_key_result(Caller, ReqId, Idx, [Bucket | Acc]);
process_keys(Caller, ReqId, Idx, {filter, Bucket, Fun}, {Bucket, K}, Acc) ->
  %% Bucket={filter,Bucket,Fun} means "only include keys
  %% in Bucket that make Fun(K) return 'true'"
  case Fun(K) of
    true ->
      buffer_key_result(Caller, ReqId, Idx, [k | Acc]);
    false ->
      Acc
  end;
process_keys(Caller, ReqId, Idx, Bucket, {Bucket, K}, Acc) ->
  buffer_key_result(Caller, ReqId, Idx, [K | Acc]);
process_keys(_Caller, _ReqId, _Idx, _Bucket, {_B, _K}, Acc) ->
  Acc.

buffer_key_result(Caller, ReqId, Idx, Acc) ->
  case length(Acc) >= 100 of
    true ->
      Caller ! {ReqId, {kl, Idx, Acc}},
      [];
    false ->
      Acc
  end.

%% @private
do_fold(Fun, Acc0, _State = #state{mod = Mod, modstate = ModState}) ->
  Mod:fold(ModState, Fun, Acc0).

%% @private
do_get_vclocks(KeyList, _State=#state{mod = Mod, modstate = ModState}) ->
  [{BKey, do_get_vclock(BKey, Mod, ModState)} || BKey <- KeyList].
%% @private
do_get_vclock(BKey, Mod, ModState) ->
  case Mod:get(ModState, BKey) of
    {error, notfound} -> vclock:fresh();
    {ok, Val} -> edynamo_object:vclock(binary_to_term(Val))
  end.

%% @private
% upon receipt of a handoff datum, there is no client FSM
do_diffobj_put(BKey={Bucket, _}, DiffObj,
               _StateData = #state{mod = Mod, modstate = ModState}) ->
  ReqID = erlang:phash2(erlang:now()),
  case syntactic_put_merge(Mod, ModState, BKey, DiffObj, ReqID) of
    {newobj, NewObj} ->
      AMObj = enforce_allow_mult(NewObj, edynamo_object:get_bucket(Bucket)),
      Val   = term_to_binary(AMObj),
      Res   = Mod:put(ModState, BKey, Val),
      case Res of
        ok -> 
          %%riak_kv_stat:update(vnode_put);
          ok;
        _ ->
          nop
      end,
      Res;
    _ ->
      ok
  end.      
